package com.welson.db.cryptor.core;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.welson.db.cryptor.annotation.CryptField;
import com.welson.db.cryptor.config.DbCryptorProperties;
import com.welson.db.cryptor.enums.CryptOperation;
import com.welson.db.cryptor.util.Cryptor;
import lombok.extern.slf4j.Slf4j;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.springframework.util.ReflectionUtils;

import java.lang.reflect.Field;
import java.util.List;
import java.util.Map;

/**
 * 实体加解密处理器工厂
 * @author welson
 */
@Slf4j
public class EntityCryptProcessorFactory {
    private static final String PROCESSOR_CLASS_SUFFIX = "$CryptProcessor";
    private static final Map<Class<?>, EntityCryptProcessor<?>> PROCESSOR_CACHE = Maps.newConcurrentMap();
    private static final EntityCryptProcessorClassLoader CLASS_LOADER = new EntityCryptProcessorClassLoader();

    private final DbCryptorProperties properties;

    public EntityCryptProcessorFactory(DbCryptorProperties properties) {
        this.properties = properties;
    }

    @SuppressWarnings("unchecked")
    public <T> EntityCryptProcessor<T> getProcessor(T entity) {
        return getProcessor((Class<T>) entity.getClass());
    }

    @SuppressWarnings("unchecked")
    public <T> EntityCryptProcessor<T> getProcessor(Class<T> clazz) {
        EntityCryptProcessor<T> processor = (EntityCryptProcessor<T>) PROCESSOR_CACHE.get(clazz);
        if (processor != null) {
            return processor;
        }

        if (!isEligibleEntityClass(clazz)) {
            PROCESSOR_CACHE.put(clazz, NoOpEntityCryptProcessor.INSTANCE);
            return (EntityCryptProcessor<T>) NoOpEntityCryptProcessor.INSTANCE;
        }

        final List<Field> validFields = Lists.newArrayList();
        ReflectionUtils.doWithFields(clazz, field -> {
            if (shouldEncrypt(field) || shouldDecrypt(field)) {
                validFields.add(field);
            }
        });

        if (validFields.isEmpty()) {
            PROCESSOR_CACHE.put(clazz, NoOpEntityCryptProcessor.INSTANCE);
            return (EntityCryptProcessor<T>) NoOpEntityCryptProcessor.INSTANCE;
        }

        try {
            String className = clazz.getName() + PROCESSOR_CLASS_SUFFIX;
            byte[] bytecode = generateProcessorClass(clazz, validFields);
            Class<?> loadedClass = CLASS_LOADER.loadClass(className, bytecode);
            processor = (EntityCryptProcessor<T>) loadedClass.newInstance();
            PROCESSOR_CACHE.put(clazz, processor);
            return processor;
        } catch (Exception e) {
            log.error("Failed to create processor for class: {}", clazz.getName(), e);
            return (EntityCryptProcessor<T>) NoOpEntityCryptProcessor.INSTANCE;
        }
    }

    private byte[] generateProcessorClass(Class<?> clazz, List<Field> validFields) {
        String className = clazz.getName();
        String bytecodeClassName = className.replace('.', '/');
        String processorClassName = className + PROCESSOR_CLASS_SUFFIX;
        String bytecodeProcessorClassName = processorClassName.replace('.', '/');
        String processorInterfaceName = EntityCryptProcessor.class.getName().replace('.', '/');

        ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
        String interfaceSignature = "Ljava/lang/Object;L" + processorInterfaceName + "<L" + bytecodeClassName + ";>;";
        cw.visit(Opcodes.V1_8, Opcodes.ACC_PUBLIC, bytecodeProcessorClassName, interfaceSignature, "java/lang/Object", new String[]{processorInterfaceName});

        // 生成构造方法
        generateConstructor(cw);

        // 生成加密方法
        generateEncryptMethod(cw, bytecodeClassName, validFields);
        generateEncryptBridgeMethod(cw, bytecodeClassName);

        // 生成解密方法
        generateDecryptMethod(cw, bytecodeClassName, validFields);
        generateDecryptBridgeMethod(cw, bytecodeClassName);

        cw.visitEnd();
        return cw.toByteArray();
    }

    private void generateConstructor(ClassVisitor cv) {
        MethodVisitor mv = cv.visitMethod(Opcodes.ACC_PUBLIC, "<init>", "()V", null, null);
        mv.visitCode();
        mv.visitVarInsn(Opcodes.ALOAD, 0);
        mv.visitMethodInsn(Opcodes.INVOKESPECIAL, "java/lang/Object", "<init>", "()V", false);
        mv.visitInsn(Opcodes.RETURN);
        mv.visitMaxs(1, 1);
        mv.visitEnd();
    }

    private void generateEncryptMethod(ClassVisitor cv, String internalClassName, List<Field> validFields) {
        String cryptorClassName = Cryptor.class.getName().replace('.', '/');
        MethodVisitor mv = cv.visitMethod(Opcodes.ACC_PUBLIC, EntityCryptProcessor.ENCRYPT_METHOD_NAME, "(L" + internalClassName + ";)V", null, null);
        mv.visitCode();

        for (Field field : validFields) {
            if (shouldEncrypt(field)) {
                String fieldName = field.getName();
                String setterDesc = getSetterDesc(field);
                mv.visitVarInsn(Opcodes.ALOAD, 1);
                mv.visitVarInsn(Opcodes.ALOAD, 1);
                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, internalClassName, getGetterName(fieldName), "()Ljava/lang/String;", false);
                mv.visitMethodInsn(Opcodes.INVOKESTATIC, cryptorClassName, "encrypt", "(Ljava/lang/String;)Ljava/lang/String;", false);
                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, internalClassName, getSetterName(fieldName), setterDesc, false);
            }
        }

        mv.visitInsn(Opcodes.RETURN);
        mv.visitMaxs(2, 2);
        mv.visitEnd();
    }

    private void generateEncryptBridgeMethod(ClassVisitor cv, String internalClassName) {
        String processorClassName = internalClassName + PROCESSOR_CLASS_SUFFIX;
        String methodSignature = "(L" + internalClassName + ";)V";
        MethodVisitor mv = cv.visitMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_BRIDGE | Opcodes.ACC_SYNTHETIC, EntityCryptProcessor.ENCRYPT_METHOD_NAME, "(Ljava/lang/Object;)V", null, null);
        mv.visitCode();
        mv.visitVarInsn(Opcodes.ALOAD, 0);
        mv.visitVarInsn(Opcodes.ALOAD, 1);
        mv.visitTypeInsn(Opcodes.CHECKCAST, internalClassName);
        mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, processorClassName, EntityCryptProcessor.ENCRYPT_METHOD_NAME, methodSignature, false);
        mv.visitInsn(Opcodes.RETURN);
        mv.visitMaxs(2, 2);
        mv.visitEnd();
    }

    private void generateDecryptMethod(ClassVisitor cv, String internalClassName, List<Field> validFields) {
        String cryptorClassName = Cryptor.class.getName().replace('.', '/');
        MethodVisitor mv = cv.visitMethod(Opcodes.ACC_PUBLIC, EntityCryptProcessor.DECRYPT_METHOD_NAME, "(L" + internalClassName + ";)V", null, null);
        mv.visitCode();

        for (Field field : validFields) {
            if (shouldDecrypt(field)) {
                String fieldName = field.getName();
                String setterDesc = getSetterDesc(field);
                mv.visitVarInsn(Opcodes.ALOAD, 1);
                mv.visitVarInsn(Opcodes.ALOAD, 1);
                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, internalClassName, getGetterName(fieldName), "()Ljava/lang/String;", false);
                mv.visitMethodInsn(Opcodes.INVOKESTATIC, cryptorClassName, "decrypt", "(Ljava/lang/String;)Ljava/lang/String;", false);
                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, internalClassName, getSetterName(fieldName), setterDesc, false);
            }
        }

        mv.visitInsn(Opcodes.RETURN);
        mv.visitMaxs(2, 2);
        mv.visitEnd();
    }

    private void generateDecryptBridgeMethod(ClassVisitor cv, String internalClassName) {
        String processorClassName = internalClassName + PROCESSOR_CLASS_SUFFIX;
        String methodSignature = "(L" + internalClassName + ";)V";
        MethodVisitor mv = cv.visitMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_BRIDGE | Opcodes.ACC_SYNTHETIC, EntityCryptProcessor.DECRYPT_METHOD_NAME, "(Ljava/lang/Object;)V", null, null);
        mv.visitCode();
        mv.visitVarInsn(Opcodes.ALOAD, 0);
        mv.visitVarInsn(Opcodes.ALOAD, 1);
        mv.visitTypeInsn(Opcodes.CHECKCAST, internalClassName);
        mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, processorClassName, EntityCryptProcessor.DECRYPT_METHOD_NAME, methodSignature, false);
        mv.visitInsn(Opcodes.RETURN);
        mv.visitMaxs(2, 2);
        mv.visitEnd();
    }

    private boolean shouldEncrypt(Field field) {
        if (field == null) {
            return false;
        }
        CryptField annotation = field.getAnnotation(CryptField.class);
        return annotation != null && (annotation.operation() == CryptOperation.ENCRYPT || annotation.operation() == CryptOperation.BOTH);
    }

    private boolean shouldDecrypt(Field field) {
        if (field == null) {
            return false;
        }
        CryptField annotation = field.getAnnotation(CryptField.class);
        return annotation != null && (annotation.operation() == CryptOperation.DECRYPT || annotation.operation() == CryptOperation.BOTH);
    }

    private String getGetterName(String fieldName) {
        return "get" + capitalize(fieldName);
    }

    private String getSetterName(String fieldName) {
        return "set" + capitalize(fieldName);
    }

    private String getSetterDesc(Field field) {
        String fieldName = field.getName();
        String setterName = getSetterName(fieldName);
        Class<?> fieldDeclaringClass = field.getDeclaringClass();
        Class<?> returnType;
        try {
            returnType = fieldDeclaringClass.getMethod(setterName, String.class).getReturnType();
        } catch (NoSuchMethodException e) {
            throw new RuntimeException("Failed to get setter method", e);
        }
        return returnType == void.class || returnType == Void.class ? "(Ljava/lang/String;)V" : "(Ljava/lang/String;)L" + returnType.getName().replace('.', '/') + ";";
    }

    private String capitalize(String str) {
        if (str == null || str.isEmpty()) {
            return str;
        }
        return str.substring(0, 1).toUpperCase() + str.substring(1);
    }

    private boolean isEligibleEntityClass(Class<?> entityClass) {
        if (entityClass == null) {
            return false;
        }
        String className = entityClass.getName();
        return properties.getScanPackages().stream().anyMatch(className::startsWith);
    }

    private static class EntityCryptProcessorClassLoader extends ClassLoader {
        public EntityCryptProcessorClassLoader() {
            super(Thread.currentThread().getContextClassLoader());
        }

        public Class<?> loadClass(String className, byte[] classBytes) {
            try {
                return loadClass(className);
            } catch (ClassNotFoundException e) {
                return defineClass(className, classBytes, 0, classBytes.length);
            }
        }
    }
} 